# -*- coding: utf-8 -*-
# @Author  : WangFeng
# @Time    : 2024/4/16 15:13
# @File    : nsw.py
# @Software: PyCharm

from __future__ import print_function, division
import random
import math
from array import array
import time


class Node(object):
    def __init__(self, idx, vector):
        self.idx = idx
        self.vector = array('d', vector)
        self.connections = []


class NSWGraph(object):
    def __init__(self, max_neighbors):
        if max_neighbors < 1:
            raise ValueError("max_neighbors must be greater than or equal to 1.")
        self.nodes = []
        self.dimension = None
        self.max_neighbors = max_neighbors

    def add_node(self, vector):
        node = Node(len(self.nodes), vector)
        self._add_node(node)

    def add_nodes(self, vectors):
        nodes = [Node(i, vector) for i, vector in enumerate(vectors)]
        random.shuffle(nodes)
        for node in nodes:
            self._add_node(node)

    def _add_node(self, node):
        # 将要添加的结点放到位图列表中
        self.nodes.append(node)
        # 初始化结点向量的维度
        if self.dimension is None:
            self.dimension = len(node.vector)

        # 如果结点列表长度小于等于最近邻的K值，那么直接连，构造快速通道 long range
        if len(self.nodes) <= self.max_neighbors:
            for existing_node in self.nodes[:-1]:
                existing_node.connections.append(node)
                node.connections.append(existing_node)
        # 如果结点列表长度大于最近邻的K值，那么需要求 离待插入结点最近的K个结点作为路径链接
        else:
            nearest_neighbors = self._find_k_nearest(node, self.max_neighbors)
            for nearest_neighbor, _ in nearest_neighbors:
                nearest_neighbor.connections.append(node)
                node.connections.append(nearest_neighbor)

    def _find_k_nearest(self, node, k):
        # 给定一个query node 近邻查找离它最近的结点
        nearest_node, _ = self.find_nearest(node)
        # 加载离query最近点的所有邻居
        candidates = nearest_node.connections.copy()
        # 初始化最近邻结点tuple列表(结点,离query的距离)
        nearest_nodes = [(nearest_node, euclidean_distance(node.vector, nearest_node.vector))]

        while len(nearest_nodes) < k * 1.8 and candidates:
            current_node = candidates.pop()
            distance = euclidean_distance(node.vector, current_node.vector)
            nearest_nodes.append((current_node, distance))
            candidates.extend(current_node.connections)

        nearest_nodes.sort(key=lambda x: x[1])
        return nearest_nodes[:k]

    # 近邻查找算法
    def find_nearest(self, node):
        if len(self.nodes) == 1:
            # 如果位图列表中只有一个结点，那么直接返回
            return self.nodes[0], euclidean_distance(node.vector, self.nodes[0].vector)
        else:
            best_distance = float('inf')
            best_node = self.nodes[0]  # 随机初始化一个结点作为query的entry_point

            while True:
                # 遍历entry_point的邻居结点，计算与query的距离
                connections = [(n, euclidean_distance(node.vector, n.vector)) for n in best_node.connections]
                # 取出entry_point的邻居结点中，离query最近的邻居结点
                current_node, current_distance = min(connections, key=lambda x: x[1])
                # 如果entry_point的邻居结点比entry_point 更接近与query，那么继续寻找
                if current_distance < best_distance:
                    best_distance = current_distance
                    best_node = current_node
                # 如果entry_point的邻居结点都比entry_point 距离query远，那么这个entry_point就是距离query最近的结点
                else:
                    break
            return best_node, best_distance


def euclidean_distance(v1, v2):
    return math.sqrt(sum((x - y) ** 2 for x, y in zip(v1, v2)))


if __name__ == '__main__':
    n = 100
    start_init_time = time.time()
    graph = NSWGraph(10)
    points = [array('d', [random.random(), random.random(), random.random()]) for _ in range(n)]
    graph.add_nodes(points)
    end_init_time = time.time()

    init_time = float(end_init_time - start_init_time)
    print("Initialization Time for {} Nodes: {} seconds".format(n, init_time))

    import json
    import pandas as pd

    ori_lst = [json.loads(i) for i in open(r'c:/Users/Administrator/Desktop/tmp_output', 'r')]
    bad_case_lst = [i.split('\t')[0] for i in open(r'c:/Users/Administrator/Desktop/debug_eval')]
    match_lst = [j for i in bad_case_lst for j in ori_lst if i == j['query']]
    res = []
    for i in match_lst:
        dic = {'query': i['query']}
        retrieval_lst = json.loads(i['retrieval'])
        dic.update({k: (str(retrieval_lst[idx]['score']) + '|' + (
            retrieval_lst[idx]['query'] if len(retrieval_lst[idx]['query']) != 0 else retrieval_lst[idx]['doc'])) for
                    idx, k in enumerate([f'向量{j}' for j in range(1, len(retrieval_lst) + 1)])})
        for idx, k in enumerate([f'向量{j}' for j in range(1, len(retrieval_lst) + 1)]):
            dic[k] = 'doc|' + dic[k] if len(retrieval_lst[idx]['id']) > 5 else 'operation|' + dic[k]
        res.append(dic)
    df = pd.DataFrame(res)
    df.to_excel('./bad_case.xlsx', index=False)
